package cn.bugstack.xfg.dev.tech.aop;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

/**
 * Redis 实现限流（改进方案：使用lua脚本）
 * 如果做排队，则应该怎么实现
 */
@Component
public class RedisRateLimiterAOP {

    private final RedisTemplate<String, String> redisTemplate;

    public RedisRateLimiterAOP(RedisTemplate<String, String> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    public boolean allowRequest(String key, int permitsPerSecond, int blacklistThreshold) {
        String rateLimiterKey = "ratelimiter:" + key;
        String blacklistKey = "blacklist:" + key;

        // 检查黑名单
        Boolean isBlacklisted = redisTemplate.hasKey(blacklistKey);
        if (Boolean.TRUE.equals(isBlacklisted)) {
            return false;
        }

        // 实现令牌桶算法
        long currentTime = System.currentTimeMillis();
        long lastRefillTime = redisTemplate.opsForValue().get(rateLimiterKey + ":last_refill") != null ?
                Long.parseLong(redisTemplate.opsForValue().get(rateLimiterKey + ":last_refill")) : 0;
        long tokens = redisTemplate.opsForValue().get(rateLimiterKey + ":tokens") != null ?
                Long.parseLong(redisTemplate.opsForValue().get(rateLimiterKey + ":tokens")) : permitsPerSecond;

        // 计算需要添加的令牌数
        long newTokens = Math.min(permitsPerSecond,
                tokens + (currentTime - lastRefillTime) / 1000 * permitsPerSecond);

        // 如果有足够的令牌，允许请求
        if (newTokens > 0) {
            redisTemplate.opsForValue().set(rateLimiterKey + ":tokens", String.valueOf(newTokens - 1));
            redisTemplate.opsForValue().set(rateLimiterKey + ":last_refill", String.valueOf(currentTime));
            return true;
        } else {
            // 增加黑名单计数
            Long blacklistCount = redisTemplate.opsForValue().increment(rateLimiterKey + ":blacklist_count");
            if (blacklistCount != null && blacklistCount > blacklistThreshold) {
                redisTemplate.opsForValue().set(blacklistKey, "1", 24, TimeUnit.HOURS);
            }
            return false;
        }
    }

    // 使用示例
    //@Autowired
    //private RedisRateLimiter redisRateLimiter;
    //
    //@Around("aopPoint() && @annotation(accessInterceptor)")
    //public Object doRouter(ProceedingJoinPoint jp, AccessInterceptor accessInterceptor) throws Throwable {
    //    String key = getAttrValue(accessInterceptor.key(), jp.getArgs());
    //
    //    if (!redisRateLimiter.allowRequest(key, accessInterceptor.permitsPerSecond(), accessInterceptor.blacklistCount())) {
    //        return fallbackMethodResult(jp, accessInterceptor.fallbackMethod());
    //    }
    //
    //    return jp.proceed();
    //}
}